import cv2
import numpy as np
from PIL import Image

# 1. 加载图片
input_path = "./ziyuan/2.png"  # 替换为实际图片路径
image = cv2.imread(input_path, cv2.IMREAD_UNCHANGED)

# 检查是否有 alpha 通道
if image.shape[2] == 4:  # 如果图片有透明通道
    alpha_channel = image[:, :, 3]
    rgb_image = image[:, :, :3]
else:
    alpha_channel = None
    rgb_image = image

# 2. 转换为 HSV 以提取背景和前景
hsv_image = cv2.cvtColor(rgb_image, cv2.COLOR_BGR2HSV)

# 3. 创建掩码（根据颜色范围调整，适合当前背景）
lower_bound = np.array([0, 0, 0])  # 黑色背景下限
upper_bound = np.array([180, 255, 30])  # 黑色背景上限
mask = cv2.inRange(hsv_image, lower_bound, upper_bound)

# 反转掩码，保留人体
foreground_mask = cv2.bitwise_not(mask)

# 4. 创建透明背景图像
output = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)

# 设置前景为白色，背景为透明
output[:, :, 0:3] = 255  # 设置前景颜色为白色
output[:, :, 3] = foreground_mask  # 使用掩码定义透明度

# 5. 保存结果
output_image = Image.fromarray(output)
output_image.save("output.png")
output_image.show()